import torch

from ..a3_netstructure import Brain
from ..a1_file_version_control_system import result_vcs
from ..a5_result_methods import get_analyse_result
from QfUtil.TrainTool.train_classifier import tell_accuracy
from .. import node_sizes
from ..a4_train_methods import get_trained_brain
from matplotlib import pyplot as plt
from ..a2_load_dataset import get_samples
from ..a5_result_methods import yield_condition_and_results_for_noises_and_brains


def accuracy_from_(result, label):
    """

    :param result:
    :param label:
    :return: va, aa, ca
    """
    vr = result[0]['layer 6']
    ar = result[1]['layer 6']
    cr = result[2]['output']

    va = tell_accuracy(vr, label)
    aa = tell_accuracy(ar, label)
    ca = tell_accuracy(cr, label)

    return va, aa, ca


def get_accuracy_of(brain: Brain, samples=None, version=None, analyse_version=None):
    """

    :param samples:
    :param brain:
    :param version:
    :return: va, aa, ca
    """
    result_name = 'vac_accuracies_d' + str(brain.depth) + '_s' + str(brain.node_size) + '.result'
    if version is not None:
        return result_vcs.load(result_name, version)

    vdin, adin, label = samples

    vr, ar, cr = get_analyse_result(samples, brain, analyse_version)

    v_predict = vr['layer 6']
    a_predict = ar['layer 6']
    c_predict = cr['output']

    va = tell_accuracy(v_predict, label)
    aa = tell_accuracy(a_predict, label)
    ca = tell_accuracy(c_predict, label)

    result = (va, aa, ca)

    result_vcs.save(result, result_name, version)

    return result


# region count accuracy differed by node_size
def count_accuracies_of_different_node_size(depth=1, node_sizes=node_sizes, version=None, brain_version=None,
                                            sample_version=None, analyse_version=None):
    """

    :param depth:
    :param node_sizes:
    :param version:
    :param brain_version:
    :return: accuracies: Tensor(3*4)
    """
    name = 'accuracies_of_different_node_size.result'
    if version is not None:
        return result_vcs.load(name, version)

    samples = get_samples(version=sample_version)
    accuracies = [[], [], []]
    for node_size in node_sizes:
        brain = get_trained_brain(depth=depth, node_size=node_size, version=brain_version)
        va, aa, ca = get_accuracy_of(samples=samples, brain=brain, version=version, analyse_version=analyse_version)
        accuracies[0].append(va), accuracies[1].append(aa), accuracies[2].append(ca)

    accuracies = torch.Tensor(accuracies)

    result_vcs.save(accuracies, name, version)

    return accuracies


def draw_accuracies(accuracies):
    """

    :param name:
    :param version: must not be None
    :return:
    """
    figure = plt.figure()
    plt.plot(accuracies.cpu().detach().T)
    return figure


# endregion

# region noise_strength-accuracy--node_size
def get_accuracy_to_noises_of_nodesizes(results_version, label_version, brains_version, version=None):
    """

    :param results_version: not None, meaning should have results before.
    :param version:
    :return:
    """

    name = 'to_noises_of_nodesizes.accuracies'
    if version is not None:
        return result_vcs.load(name, version)

    sample = get_samples(version=label_version)
    label = sample[2]

    va_accuracies = {}
    aa_accuracies = {}
    ca_accuracies = {}
    for noise_strength, node_size, result in yield_condition_and_results_for_noises_and_brains(
            version=results_version,
            sample_version=label_version,
            brains_version=brains_version,
    ):
        va, aa, ca = accuracy_from_(result, label)
        if node_size not in va_accuracies:
            va_accuracies[node_size] = []
        va_accuracies[node_size].append(va)
        if node_size not in aa_accuracies:
            aa_accuracies[node_size] = []
        aa_accuracies[node_size].append(aa)
        if node_size not in ca_accuracies:
            ca_accuracies[node_size] = []
        ca_accuracies[node_size].append(ca)

    result = (va_accuracies, aa_accuracies, ca_accuracies)
    result_vcs.save(result, name)

    return va_accuracies, aa_accuracies, ca_accuracies


# endregion

# region accuracy of different digit
def accuracy_of_digits_from_(result, label):
    va = []
    aa = []
    ca = []

    vr = result[0]['layer 6']
    ar = result[1]['layer 6']
    cr = result[2]['output']

    for i in range(10):
        va.append(tell_accuracy(vr[label == i], label[label == i]))
        aa.append(tell_accuracy(ar[label == i], label[label == i]))
        ca.append(tell_accuracy(cr[label == i], label[label == i]))

    return va, aa, ca

# region noise_strength-accuracy--node_size
def get_accuracy_of_digits_to_noises_of_nodesizes(results_version, label_version, brains_version, version=None):
    """

    :param results_version: not None, meaning should have results before.
    :param version:
    :return:
    """

    name = 'of_digits_to_noises_of_nodesizes.accuracies'
    if version is not None:
        return result_vcs.load(name, version)

    sample = get_samples(version=label_version)
    label = sample[2]

    va_accuracies = {}
    aa_accuracies = {}
    ca_accuracies = {}
    for noise_strength, node_size, result in yield_condition_and_results_for_noises_and_brains(
            version=results_version,
            sample_version=label_version,
            brains_version=brains_version,
    ):
        va, aa, ca = accuracy_of_digits_from_(result, label)
        if node_size not in va_accuracies:
            va_accuracies[node_size] = []
        va_accuracies[node_size].append(va)
        if node_size not in aa_accuracies:
            aa_accuracies[node_size] = []
        aa_accuracies[node_size].append(aa)
        if node_size not in ca_accuracies:
            ca_accuracies[node_size] = []
        ca_accuracies[node_size].append(ca)

    result = (va_accuracies, aa_accuracies, ca_accuracies)
    result_vcs.save(result, name)

    return va_accuracies, aa_accuracies, ca_accuracies


# endregion
# endregion
